{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "import json\n",
    "from pathlib import Path\n",
    "\n",
    "from virtual_lab.constants import CONSISTENT_TEMPERATURE, CREATIVE_TEMPERATURE\n",
    "from virtual_lab.prompts import (\n",
    "    CODING_RULES,\n",
    "    REWRITE_PROMPT,\n",
    "    create_merge_prompt,\n",
    ")\n",
    "from virtual_lab.run_meeting import run_meeting\n",
    "from virtual_lab.utils import load_summaries\n",
    "\n",
    "from nanobody_constants import (\n",
    "    background_prompt,\n",
    "    nanobody_prompt,\n",
    "    num_iterations,\n",
    "    num_rounds,\n",
    "    discussions_phase_to_dir,\n",
    "    principal_investigator,\n",
    "    team_members,\n",
    "    machine_learning_specialist,\n",
    "    computational_biologist,\n",
    "    scientific_critic\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9259ddc2710488f",
   "metadata": {},
   "source": [
    "## Team selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c55d79aa92e377b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Team selection - prompts\n",
    "team_selection_agenda = f\"\"\"{background_prompt} You need to select a team of three scientists to help you with this project. Please select the team members that you would like to invite to a discussion to create the antibody/nanobody design approach. Please list the team members in the following format, using the team member below as an example. You should not include yourself (Principal Investigator) in the list.\n",
    "\n",
    "Agent(\n",
    "    title=\"Principal Investigator\",\n",
    "    expertise=\"applying artificial intelligence to biomedical research\",\n",
    "    goal=\"perform research in your area of expertise that maximizes the scientific impact of the work\",\n",
    "    role=\"lead a team of experts to solve an important problem in artificial intelligence for biomedicine, make key decisions about the project direction based on team member input, and manage the project timeline and resources\",\n",
    ")\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50722e4dadc135d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Team selection - discussion\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    concurrent.futures.wait([\n",
    "        executor.submit(\n",
    "            run_meeting,\n",
    "            meeting_type=\"individual\",\n",
    "            team_member=principal_investigator,\n",
    "            agenda=team_selection_agenda,\n",
    "            save_dir=discussions_phase_to_dir[\"team_selection\"],\n",
    "            save_name=f\"discussion_{iteration_num + 1}\",\n",
    "            temperature=CREATIVE_TEMPERATURE,\n",
    "        ) for iteration_num in range(num_iterations)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1c023cebeaaf883",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Team selection - merge\n",
    "team_selection_summaries = load_summaries(\n",
    "    discussion_paths=list(discussions_phase_to_dir[\"team_selection\"].glob(\"discussion_*.json\")))\n",
    "print(f\"Number of summaries: {len(team_selection_summaries)}\")\n",
    "\n",
    "team_selection_merge_prompt = create_merge_prompt(agenda=team_selection_agenda)\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=principal_investigator,\n",
    "    summaries=team_selection_summaries,\n",
    "    agenda=team_selection_merge_prompt,\n",
    "    save_dir=discussions_phase_to_dir[\"team_selection\"],\n",
    "    save_name=\"merged\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3804141d0b26537",
   "metadata": {},
   "source": [
    "## Projects specification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9df84a9e45d31bbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Project specification - prompts\n",
    "project_specification_agenda = f\"{background_prompt} Please create an antibody/nanobody design approach to solve this problem. Decide whether you will design antibodies or nanobodies. For your choice, decide whether you will design the antibodies/nanobodies de novo or whether you will modify existing antibodies/nanobodies. If modifying existing antibodies/nanobodies, please specify which antibodies/nanobodies to start with as good candidates for targeting the newest variant of the SARS-CoV-2 spike protein. If designing antibodies/nanobodies de novo, please describe how you will propose antibody/nanobody candidates.\"\n",
    "\n",
    "project_specification_questions = (\n",
    "    \"Will you design standard antibodies or nanobodies?\",\n",
    "    \"Will you design antibodies/nanobodies de novo or will you modify existing antibodies/nanobodies (choose only one)?\",\n",
    "    \"If modifying existing antibodies/nanobodies, which precise antibodies/nanobodies will you modify (please list 3-4)?\",\n",
    "    \"If designing antibodies/nanobodies de novo, how exactly will you propose antibody/nanobody candidates?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ff509ba8824b8ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Project specification - discussion\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    concurrent.futures.wait([\n",
    "        executor.submit(\n",
    "            run_meeting,\n",
    "            meeting_type=\"team\",\n",
    "            team_lead=principal_investigator,\n",
    "            team_members=team_members,\n",
    "            agenda=project_specification_agenda,\n",
    "            agenda_questions=project_specification_questions,\n",
    "            save_dir=discussions_phase_to_dir[\"project_specification\"],\n",
    "            save_name=f\"discussion_{iteration_num + 1}\",\n",
    "            temperature=CREATIVE_TEMPERATURE,\n",
    "            num_rounds=num_rounds,\n",
    "        ) for iteration_num in range(num_iterations)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e66e0cfa85f2176d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Project specification - merge\n",
    "project_specification_summaries = load_summaries(\n",
    "    discussion_paths=list(discussions_phase_to_dir[\"project_specification\"].glob(\"discussion_*.json\")))\n",
    "print(f\"Number of summaries: {len(project_specification_summaries)}\")\n",
    "\n",
    "project_specification_merge_prompt = create_merge_prompt(\n",
    "    agenda=project_specification_agenda,\n",
    "    agenda_questions=project_specification_questions,\n",
    ")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=principal_investigator,\n",
    "    summaries=project_specification_summaries,\n",
    "    agenda=project_specification_merge_prompt,\n",
    "    save_dir=discussions_phase_to_dir[\"project_specification\"],\n",
    "    save_name=\"merged\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    "    num_rounds=num_rounds,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e5307d354d28011",
   "metadata": {},
   "source": [
    "## Tools Selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a001a2ca1ea13a4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tools selection - prompts\n",
    "tools_selection_agenda = f\"{background_prompt} {nanobody_prompt} Now you need to select machine learning and/or computational tools to implement this nanobody design approach. Please list several tools (5-10) that would be relevant to this nanobody design approach and how they could be used in the context of this project. If selecting machine learning tools, please prioritize pre-trained models (e.g., pre-trained protein language models or protein structure prediction models) for simplicity.\"\n",
    "\n",
    "tools_selection_questions = (\n",
    "    \"What machine learning and/or computational tools could be used for this nanobody design approach (list 5-10)?\",\n",
    "    \"For each tool, how could it be used for designing modified nanobodies?\",\n",
    ")\n",
    "\n",
    "tools_selection_prior_summaries = load_summaries(\n",
    "    discussion_paths=[discussions_phase_to_dir[\"project_specification\"] / \"merged.json\"])\n",
    "print(f\"Number of prior summaries: {len(tools_selection_prior_summaries)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a95de22d2f1d277b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tools selection - discussion\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    concurrent.futures.wait([\n",
    "        executor.submit(\n",
    "            run_meeting,\n",
    "            meeting_type=\"team\",\n",
    "            team_lead=principal_investigator,\n",
    "            team_members=team_members,\n",
    "            summaries=tools_selection_prior_summaries,\n",
    "            agenda=tools_selection_agenda,\n",
    "            agenda_questions=tools_selection_questions,\n",
    "            save_dir=discussions_phase_to_dir[\"tools_selection\"],\n",
    "            save_name=f\"discussion_{iteration_num + 1}\",\n",
    "            temperature=CREATIVE_TEMPERATURE,\n",
    "            num_rounds=num_rounds,\n",
    "        ) for iteration_num in range(num_iterations)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "580378a119b0cce5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tools selection - merge\n",
    "tools_selection_summaries = load_summaries(\n",
    "    discussion_paths=list(discussions_phase_to_dir[\"tools_selection\"].glob(\"discussion_*.json\")))\n",
    "print(f\"Number of summaries: {len(tools_selection_summaries)}\")\n",
    "\n",
    "tools_selection_merge_prompt = create_merge_prompt(\n",
    "    agenda=tools_selection_agenda,\n",
    "    agenda_questions=tools_selection_questions,\n",
    ")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=principal_investigator,\n",
    "    summaries=tools_selection_summaries,\n",
    "    agenda=tools_selection_merge_prompt,\n",
    "    save_dir=discussions_phase_to_dir[\"tools_selection\"],\n",
    "    save_name=\"merged\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    "    num_rounds=num_rounds,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3777aacd05e6e17",
   "metadata": {},
   "source": [
    "## Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4b45d33467d1405",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implementation - prompts\n",
    "implementation_agent_selection_agenda = f\"{background_prompt} {nanobody_prompt} Your team needs to build three components of a nanobody design pipeline: ESM, AlphaFold-Multimer, and Rosetta. For each component, please select the team member who will implement the component. A team member may implement more than one component.\"\n",
    "\n",
    "implementation_agent_selection_questions = (\n",
    "    \"Which team member will implement ESM?\",\n",
    "    \"Which team member will implement AlphaFold-Multimer?\",\n",
    "    \"Which team member will implement Rosetta?\",\n",
    ")\n",
    "\n",
    "implementation_agent_selection_prior_summaries = load_summaries(\n",
    "    discussion_paths=[discussions_phase_to_dir[\"team_selection\"] / \"merged.json\",\n",
    "                      discussions_phase_to_dir[\"project_specification\"] / \"merged.json\",\n",
    "                      discussions_phase_to_dir[\"tools_selection\"] / \"merged.json\"])\n",
    "print(f\"Number of prior summaries: {len(implementation_agent_selection_prior_summaries)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "777c7f17e1853145",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implementation - discussion\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    concurrent.futures.wait([\n",
    "        executor.submit(\n",
    "            run_meeting,\n",
    "            meeting_type=\"individual\",\n",
    "            team_member=principal_investigator,\n",
    "            summaries=implementation_agent_selection_prior_summaries,\n",
    "            agenda=implementation_agent_selection_agenda,\n",
    "            agenda_questions=implementation_agent_selection_questions,\n",
    "            save_dir=discussions_phase_to_dir[\"implementation_agent_selection\"],\n",
    "            save_name=f\"discussion_{iteration_num + 1}\",\n",
    "            temperature=CREATIVE_TEMPERATURE,\n",
    "        ) for iteration_num in range(num_iterations)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b741ab035e21003c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implementation - merge\n",
    "implementation_agent_selection_summaries = load_summaries(\n",
    "    discussion_paths=list(discussions_phase_to_dir[\"implementation_agent_selection\"].glob(\"discussion_*.json\")))\n",
    "print(f\"Number of summaries: {len(implementation_agent_selection_summaries)}\")\n",
    "\n",
    "implementation_agent_selection_merge_prompt = create_merge_prompt(\n",
    "    agenda=implementation_agent_selection_agenda,\n",
    "    agenda_questions=implementation_agent_selection_questions\n",
    ")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=principal_investigator,\n",
    "    summaries=implementation_agent_selection_summaries,\n",
    "    agenda=implementation_agent_selection_merge_prompt,\n",
    "    save_dir=discussions_phase_to_dir[\"implementation_agent_selection\"],\n",
    "    save_name=\"merged\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b1ac62cf9a8f39",
   "metadata": {},
   "source": [
    "### ESM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b2b6e3a180a548e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ESM - prompts\n",
    "esm_agenda = f\"{background_prompt} {nanobody_prompt} Now you must use ESM to suggest modifications to an existing antibody. Please write a complete Python script that takes a nanobody sequence as input and uses ESM amino acid log-likelihoods to identify the most promising point mutations by log-likelihood ratio.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0affb5cd00bec96",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ESM - discussion\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    concurrent.futures.wait([\n",
    "        executor.submit(\n",
    "            run_meeting,\n",
    "            meeting_type=\"individual\",\n",
    "            team_member=machine_learning_specialist,\n",
    "            agenda=esm_agenda,\n",
    "            agenda_rules=CODING_RULES,\n",
    "            save_dir=discussions_phase_to_dir[\"esm\"],\n",
    "            save_name=f\"discussion_{iteration_num + 1}\",\n",
    "            temperature=CREATIVE_TEMPERATURE,\n",
    "            num_rounds=num_rounds,\n",
    "        ) for iteration_num in range(num_iterations)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b0936e225dcc32f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ESM - merge\n",
    "esm_summaries = load_summaries(discussion_paths=list(discussions_phase_to_dir[\"esm\"].glob(\"discussion_*.json\")))\n",
    "print(f\"Number of summaries: {len(esm_summaries)}\")\n",
    "\n",
    "esm_merge_prompt = create_merge_prompt(\n",
    "    agenda=esm_agenda,\n",
    "    agenda_rules=CODING_RULES,\n",
    ")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=machine_learning_specialist,\n",
    "    summaries=esm_summaries,\n",
    "    agenda=esm_merge_prompt,\n",
    "    save_dir=discussions_phase_to_dir[\"esm\"],\n",
    "    save_name=\"merged\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cf7a7a3d2b7b7fe",
   "metadata": {},
   "source": [
    "### Improve ESM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b062e9cb5f92a482",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Improve ESM - prompts\n",
    "improve_esm_agenda = f\"\"\"You previously wrote a Python script that uses ESM to compute the log-likelihood ratio of point mutations in a nanobody sequence (see summary). {REWRITE_PROMPT}\n",
    "\n",
    "1. Replace \"facebook/esm1b-t33_650M_UR50S\" with \"facebook/esm1b_t33_650M_UR50S\".\n",
    "2. Run the calculations of the mutant log-likelihoods by iterating through the sequences in batches of 16.\n",
    "3. Add a progress bar to the batched mutant log-likelihood calculations.\n",
    "4. Run the mutant log-likelihood calculations on CUDA but with no gradients.\n",
    "5. Load the nanobody sequence from a user-specified CSV file that has the columns \"sequence\" and \"name\". Adapt your code to run the mutant log-likelihood calculations on all sequences in the CSV file one-by-one.\n",
    "6. For each sequence, save the mutant log-likelihoods to a CSV file with the format \"mutated_sequence,position,original_aa,mutated_aa,log_likelihood_ratio\". Ask the user for a save directory and then save this CSV file in that directory with the name: <nanbody-name>.csv.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6388c95f9e77d811",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Improve ESM - discussion\n",
    "improve_esm_summaries = load_summaries(discussion_paths=[discussions_phase_to_dir[\"esm\"] / \"merged.json\"])\n",
    "print(f\"Number of summaries: {len(improve_esm_summaries)}\")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=machine_learning_specialist,\n",
    "    summaries=improve_esm_summaries,\n",
    "    agenda=improve_esm_agenda,\n",
    "    save_dir=discussions_phase_to_dir[\"esm\"],\n",
    "    save_name=\"improved\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ba22302ea18a575",
   "metadata": {},
   "source": [
    "### AlphaFold-Multimer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9588993af806c3f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# AlphaFold-Multimer - prompts\n",
    "alphafold_agenda = f\"{background_prompt} {nanobody_prompt} Now you must use AlphaFold-Multimer to predict the structure of a nanobody-antigen complex and evaluate its binding. I will run AlphaFold-Multimer on several nanobody-antigen complexes and you need to process the outputs. Please write a complete Python script that takes as input a directory containing PDB files where each PDB file contains one nanobody-antigen complex predicted by AlphaFold-Multimer and outputs a CSV file containing the AlphaFold-Multimer confidence of each nanobody-antigen complex in terms of the interface pLDDT.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ca49b4e29ac8b6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# AlphaFold-Multimer - discussion\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    concurrent.futures.wait([\n",
    "        executor.submit(\n",
    "            run_meeting,\n",
    "            meeting_type=\"individual\",\n",
    "            team_member=computational_biologist,\n",
    "            agenda=alphafold_agenda,\n",
    "            agenda_rules=CODING_RULES,\n",
    "            save_dir=discussions_phase_to_dir[\"alphafold\"],\n",
    "            save_name=f\"discussion_{iteration_num + 1}\",\n",
    "            temperature=CREATIVE_TEMPERATURE,\n",
    "            num_rounds=num_rounds,\n",
    "        ) for iteration_num in range(num_iterations)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d03e25b1ecad9557",
   "metadata": {},
   "outputs": [],
   "source": [
    "# AlphaFold-Multimer - merge\n",
    "alphafold_summaries = load_summaries(\n",
    "    discussion_paths=list(discussions_phase_to_dir[\"alphafold\"].glob(\"discussion_*.json\")))\n",
    "print(f\"Number of summaries: {len(alphafold_summaries)}\")\n",
    "\n",
    "alphafold_merge_prompt = create_merge_prompt(\n",
    "    agenda=alphafold_agenda,\n",
    "    agenda_rules=CODING_RULES,\n",
    ")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=computational_biologist,\n",
    "    summaries=alphafold_summaries,\n",
    "    agenda=alphafold_merge_prompt,\n",
    "    save_dir=discussions_phase_to_dir[\"alphafold\"],\n",
    "    save_name=\"merged\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "619858f6ff3335c6",
   "metadata": {},
   "source": [
    "### Improve AlphaFold-Multimer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd78a537e8608abb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Improve AlphaFold-Multimer - prompts\n",
    "improve_alphafold_agenda = f\"\"\"You previously wrote a Python script that processes the outputs of AlphaFold-Multimer to calculate the confidence of nanobody-antigen complexes (see summary). {REWRITE_PROMPT}\n",
    "\n",
    "1. Replace the current imports of Chain and Residue with \"from Bio.PDB.Chain import Chain\" and \"from Bio.PDB.Residue import Residue\".\n",
    "2. Remove the logging setup and simply print any log messages to the console.\n",
    "3. Replace the parallel processing with sequential processing to avoid getting an \"OSError: Too many open files\".\n",
    "4. Change the list of pdb_files to instead get all PDB files in the directory that follow the pattern \"**/*unrelaxed_rank_001*.pdb\".\n",
    "5. Change the calculation of average pLDDT to divide by the number of atoms rather than the number of residues.\n",
    "6. Return and save in the CSV both the number of residues and the number of atoms in the interface.\n",
    "7. Change the default distance threshold to 4.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d9bf62680e027cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Improve AlphaFold-Multimer - discussion\n",
    "improve_alphafold_summaries = load_summaries(discussion_paths=[discussions_phase_to_dir[\"alphafold\"] / \"merged.json\"])\n",
    "print(f\"Number of summaries: {len(improve_alphafold_summaries)}\")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=computational_biologist,\n",
    "    summaries=improve_alphafold_summaries,\n",
    "    agenda=improve_alphafold_agenda,\n",
    "    save_dir=discussions_phase_to_dir[\"alphafold\"],\n",
    "    save_name=\"improved\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ba34f13b0ec3889",
   "metadata": {},
   "source": [
    "### Rosetta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38b7c084b07dd169",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rosetta - prompts\n",
    "rosetta_agenda = f\"{background_prompt} {nanobody_prompt} Now you must use Rosetta to calculate the binding energy of nanobody-antigen complexes. You must do this in three parts. First, write a complete RosettaScripts XML file that calculates the binding energy of a nanobody-antigen complex as provided in PDB format, including any necessary preprocessing steps for the complex. Second, write an example command that uses Rosetta to run this RosettaScripts XML file on a given PDB file to calculate the binding energy and save it to a score file. Third, write a complete Python script that takes as input a directory with multiple Rosetta binding energy score files and outputs a single CSV file with the names and scores of each of the individual files in sorted order (highest to lowest score).\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3433242b9472cd28",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rosetta - discussion\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    concurrent.futures.wait([\n",
    "        executor.submit(\n",
    "            run_meeting,\n",
    "            meeting_type=\"individual\",\n",
    "            team_member=computational_biologist,\n",
    "            agenda=rosetta_agenda,\n",
    "            agenda_rules=CODING_RULES,\n",
    "            save_dir=discussions_phase_to_dir[\"rosetta\"],\n",
    "            save_name=f\"discussion_{iteration_num + 1}\",\n",
    "            temperature=CREATIVE_TEMPERATURE,\n",
    "            num_rounds=num_rounds,\n",
    "        ) for iteration_num in range(num_iterations)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4563b6d0b6116f2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rosetta - merge\n",
    "rosetta_summaries = load_summaries(discussion_paths=list(discussions_phase_to_dir[\"rosetta\"].glob(\"discussion_*.json\")))\n",
    "print(f\"Number of summaries: {len(rosetta_summaries)}\")\n",
    "\n",
    "rosetta_merge_prompt = create_merge_prompt(\n",
    "    agenda=rosetta_agenda,\n",
    "    agenda_rules=CODING_RULES,\n",
    ")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=computational_biologist,\n",
    "    summaries=rosetta_summaries,\n",
    "    agenda=rosetta_merge_prompt,\n",
    "    save_dir=discussions_phase_to_dir[\"rosetta\"],\n",
    "    save_name=\"merged\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4771524f6b0c8270",
   "metadata": {},
   "source": [
    "### Improve Rosetta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a03ce8ff3dc2c718",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Improve Rosetta XML - prompts\n",
    "improve_rosetta_xml_agenda = f\"\"\"You previously wrote a RosettaScripts XML file to calculate the binding affinity of a nanobody-antigen complex (see summary). {REWRITE_PROMPT}\n",
    "\n",
    "1. Replace \"ref15.wts\" with \"ref2015.wts\".\n",
    "2. Remove the InterfaceEnergy filter since it is not valid in Rosetta.\n",
    "3. Replace the entire output tag (including any nested tags) with <OUTPUT scorefxn=\"ref15\"/>.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d92fb27857e97f52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Improve Rosetta XML - discussion\n",
    "improve_rosetta_xml_summaries = load_summaries(discussion_paths=[discussions_phase_to_dir[\"rosetta\"] / \"merged.json\"])\n",
    "print(f\"Number of summaries: {len(improve_rosetta_xml_summaries)}\")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=computational_biologist,\n",
    "    summaries=improve_rosetta_xml_summaries,\n",
    "    agenda=improve_rosetta_xml_agenda,\n",
    "    save_dir=discussions_phase_to_dir[\"rosetta\"],\n",
    "    save_name=\"improved_xml\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48ed12daeae0cf1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Improve Rosetta Python - prompts\n",
    "improve_rosetta_python_agenda = f\"\"\"You previously wrote a Python script to aggregate multiple Rosetta binding energy score files into one CSV file (see summary). {REWRITE_PROMPT}\n",
    "\n",
    "1. Modify the extract_scores_from_file function so that it extracts the dG_separated value from a file of the following form.\n",
    "\n",
    "SEQUENCE:\n",
    "SCORE: total_score complex_normalized           dG_cross dG_cross/dSASAx100 dG_separated dG_separated/dSASAx100 dSASA_hphobic dSASA_int dSASA_polar delta_unsatHbonds dslf_fa13    fa_atr    fa_dun   fa_elec fa_intra_rep fa_intra_sol_xover4              fa_rep              fa_sol hbond_E_fraction hbond_bb_sc hbond_lr_bb    hbond_sc hbond_sr_bb hbonds_int lk_ball_wtd    nres_all    nres_int       omega     p_aa_pp    packstat per_residue_energy_int pro_close rama_prepro         ref    sc_value side1_normalized side1_score side2_normalized side2_score yhh_planarity description\n",
    "SCORE:    -990.807             -2.914            -21.436             -1.857      -21.436                 -1.857       774.274  1154.088     379.813            12.000    -3.867 -1928.622   376.416  -541.777        3.745              54.944             265.303            1052.322            0.053     -84.023    -130.532     -54.069     -46.266      1.000     -41.725     340.000      55.000      39.977     -81.331       0.000                 -2.699     2.349      -6.870     131.513       0.000           -2.236     -51.431           -3.031     -97.008         1.706 KP3_Ty1-G59Y_unrelaxed_rank_001_alphafold2_multimer_v3_model_3_seed_000_0001\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea4287f833216ee9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Improve Rosetta Python - discussion\n",
    "improve_rosetta_python_summaries = load_summaries(\n",
    "    discussion_paths=[discussions_phase_to_dir[\"rosetta\"] / \"merged.json\"])\n",
    "print(f\"Number of summaries: {len(improve_rosetta_python_summaries)}\")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=computational_biologist,\n",
    "    summaries=improve_rosetta_python_summaries,\n",
    "    agenda=improve_rosetta_python_agenda,\n",
    "    save_dir=discussions_phase_to_dir[\"rosetta\"],\n",
    "    save_name=\"improved_python\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a050fffc08866b10",
   "metadata": {},
   "source": [
    "## Workflow Design"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c55ec9669eeee9fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Workflow design - prompts\n",
    "workflow_design_agenda = f\"{background_prompt} {nanobody_prompt} Your team has built three components of a nanobody design pipeline: ESM, AlphaFold-Multimer, and Rosetta. Each of these three tools can be used to score a nanobody mutation based on how the mutation affects binding to an antigen. Your goal is to start with an existing nanobody and iteratively add mutations to it to improve its binding to the newest variant of the SARS-CoV-2 spike protein, resulting in 24 modified nanobodies with one or more mutations. Please determine how to use ESM, AlphaFold-Multimer, and Rosetta in this iterative design process. An important constraint is that ESM can evaluate all potential mutations to a nanobody in 5 minutes while AlphaFold-Multimer takes 30 minutes per mutation and Rosetta takes five minutes per mutation. The whole iterative process should take no more than a few days to complete. Note that AlphaFold-Multimer must be run before Rosetta on each mutation since Rosetta requires the nanobody-antigen structure predicted by AlphaFold-Multimer. Additionally, note that ESM log-likelihood ratios are generally in the range of 5-10 (higher is better), AlphaFold-Multimer interface pLDDT confidence scores are generally in the range of 60-80 (higher is better), and Rosetta binding energy scores are generally in the range of -20 to -40 (lower is better).\"\n",
    "\n",
    "workflow_design_questions = (\n",
    "    \"In each iteration, what is the order of operations for evaluating mutations with ESM, AlphaFold-Multimer, and Rosetta?\",\n",
    "    \"In each iteration, how many mutations (give a single number) will you evaluate with ESM, AlphaFold-Multimer, and Rosetta?\",\n",
    "    \"At the end of each iteration, how will you weigh the ESM, AlphaFold-Multimer, and/or Rosetta scores (give a formula) to rank the nanobody mutations?\",\n",
    "    \"At the end of each iteration, how many of the top-ranked mutations (give a single number) will you keep for the next round?\",\n",
    "    \"How will you decide how many iterations of mutations to run?\",\n",
    "    \"After all of the iterations are complete, how exactly (step-by-step) will you select the final set of 24 modified nanobodies from across the iterations for experimental validation?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3210bdc6d68bd0bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Workflow design - discussion\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    concurrent.futures.wait([\n",
    "        executor.submit(\n",
    "            run_meeting,\n",
    "            meeting_type=\"individual\",\n",
    "            team_member=principal_investigator,\n",
    "            agenda=workflow_design_agenda,\n",
    "            agenda_questions=workflow_design_questions,\n",
    "            save_dir=discussions_phase_to_dir[\"workflow_design\"],\n",
    "            save_name=f\"discussion_{iteration_num + 1}\",\n",
    "            temperature=CREATIVE_TEMPERATURE,\n",
    "        ) for iteration_num in range(num_iterations)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb4534d4913ecf88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Workflow design - merge\n",
    "workflow_design_summaries = load_summaries(\n",
    "    discussion_paths=list(discussions_phase_to_dir[\"workflow_design\"].glob(\"discussion_*.json\")))\n",
    "print(f\"Number of summaries: {len(workflow_design_summaries)}\")\n",
    "\n",
    "workflow_design_merge_prompt = create_merge_prompt(\n",
    "    agenda=workflow_design_agenda,\n",
    "    agenda_questions=workflow_design_questions,\n",
    ")\n",
    "\n",
    "run_meeting(\n",
    "    meeting_type=\"individual\",\n",
    "    team_member=principal_investigator,\n",
    "    summaries=workflow_design_summaries,\n",
    "    agenda=workflow_design_merge_prompt,\n",
    "    save_dir=discussions_phase_to_dir[\"workflow_design\"],\n",
    "    save_name=\"merged\",\n",
    "    temperature=CONSISTENT_TEMPERATURE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb5b2c7839b930aa",
   "metadata": {},
   "source": [
    "## Virtual Lab Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "74d3a3d37081ed02",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "matplotlib.rcParams.update({'font.size': 26})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "8bebeb39ba3c3b0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "figure_dir = Path(\"figures/virtual_lab_analysis\")\n",
    "figure_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "phase_to_agent_to_word_count = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "3577d08c72a2aeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Count words that the human user wrote\n",
    "phase_to_human_words = {\n",
    "    \"team_selection\": [\n",
    "        background_prompt,\n",
    "        principal_investigator.prompt,\n",
    "        scientific_critic.prompt,\n",
    "        team_selection_agenda.replace(f\"{background_prompt} \", \"\"),\n",
    "    ],\n",
    "    \"project_specification\": [\n",
    "        project_specification_agenda.replace(f\"{background_prompt} \", \"\"),\n",
    "        *project_specification_questions,\n",
    "        nanobody_prompt,\n",
    "    ],\n",
    "    \"tools_selection\": [\n",
    "        tools_selection_agenda.replace(f\"{background_prompt} {nanobody_prompt} \", \"\"),\n",
    "        *tools_selection_questions,\n",
    "    ],\n",
    "    \"implementation_agent_selection\": [\n",
    "        implementation_agent_selection_agenda.replace(f\"{background_prompt} {nanobody_prompt} \", \"\"),\n",
    "        *implementation_agent_selection_questions,\n",
    "    ],\n",
    "    \"esm\": [\n",
    "        esm_agenda.replace(f\"{background_prompt} {nanobody_prompt} \", \"\"),\n",
    "        improve_esm_agenda.replace(f\" {REWRITE_PROMPT}\", \"\"),\n",
    "    ],\n",
    "    \"alphafold\": [\n",
    "        alphafold_agenda.replace(f\"{background_prompt} {nanobody_prompt} \", \"\"),\n",
    "        improve_alphafold_agenda.replace(f\" {REWRITE_PROMPT}\", \"\"),\n",
    "    ],\n",
    "    \"rosetta\": [\n",
    "        rosetta_agenda.replace(f\"{background_prompt} {nanobody_prompt} \", \"\"),\n",
    "        improve_rosetta_xml_agenda.replace(f\" {REWRITE_PROMPT}\", \"\"),\n",
    "        improve_rosetta_python_agenda.replace(f\" {REWRITE_PROMPT}\", \"\"),\n",
    "    ],\n",
    "    \"workflow_design\": [\n",
    "        workflow_design_agenda.replace(f\"{background_prompt} {nanobody_prompt} \", \"\"),\n",
    "        *workflow_design_questions,\n",
    "    ],\n",
    "}\n",
    "\n",
    "for phase, human_words in phase_to_human_words.items():\n",
    "    phase_to_agent_to_word_count[phase] = {\"Human Researcher\": len(\" \".join(human_words).split())}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30b622f7378bd3ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Count words that the LLM agents wrote\n",
    "for phase_name in [\"team_selection\", \"project_specification\", \"tools_selection\",\n",
    "                   \"implementation_agent_selection\", \"esm\", \"alphafold\", \"rosetta\", \"workflow_design\"]:\n",
    "    phase_dir = discussions_phase_to_dir[phase_name]\n",
    "\n",
    "    print(f\"Phase: {phase_name}\")\n",
    "\n",
    "    # Load the text written by each agent\n",
    "    agent_to_text = {}\n",
    "    for path in phase_dir.glob(\"*.json\"):\n",
    "        with open(path) as f:\n",
    "            discussion = json.load(f)\n",
    "\n",
    "        for message in discussion:\n",
    "            agent_to_text.setdefault(message[\"agent\"], []).append(message[\"message\"])\n",
    "\n",
    "    # Count the number of words written by each agent\n",
    "    for agent, text in agent_to_text.items():\n",
    "        if agent == \"User\":\n",
    "            continue\n",
    "\n",
    "        agent_to_text[agent] = \" \".join(text)\n",
    "        word_count = len(agent_to_text[agent].split())\n",
    "        phase_to_agent_to_word_count[phase_name][agent] = word_count\n",
    "\n",
    "# Print words by phase\n",
    "for phase in phase_to_agent_to_word_count:\n",
    "    print(f\"Phase: {phase}\")\n",
    "    for agent, word_count in phase_to_agent_to_word_count[phase].items():\n",
    "        print(f\"Number of words written by {agent}: {word_count:,}\")\n",
    "    print()\n",
    "\n",
    "# Sum word counts across phases\n",
    "agent_to_word_count = {}\n",
    "for phase in phase_to_agent_to_word_count:\n",
    "    for agent, word_count in phase_to_agent_to_word_count[phase].items():\n",
    "        agent_to_word_count[agent] = agent_to_word_count.get(agent, 0) + word_count\n",
    "\n",
    "# Total number of words written by each LLM agent\n",
    "for agent, word_count in agent_to_word_count.items():\n",
    "    print(f\"Total number of words written by {agent}: {word_count:,}\")\n",
    "\n",
    "print()\n",
    "\n",
    "# Total number of words written by all LLM agents\n",
    "total_human_words = sum(\n",
    "    phase_to_agent_to_word_count[phase][\"Human Researcher\"] for phase in phase_to_agent_to_word_count)\n",
    "total_agent_words = sum(word_count for agent, word_count in agent_to_word_count.items() if agent != \"Human Researcher\")\n",
    "\n",
    "print(f\"Total number of words written by Human Researcher: {total_human_words:,}\")\n",
    "print(f\"Total number of words written by all LLM agents: {total_agent_words:,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "62dd0609f9e7274c",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_to_color = {\n",
    "    agent: sns.color_palette(\"tab10\", n_colors=len(agent_to_word_count))[i]\n",
    "    for i, agent in enumerate(agent_to_word_count)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a00e6cdaa1a1c05",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8, 6))\n",
    "ax.pie(\n",
    "    agent_to_word_count.values(),\n",
    "    labels=agent_to_word_count.keys(),\n",
    "    autopct=\"%1.1f%%\",\n",
    "    colors=[agent_to_color[agent] for agent in agent_to_word_count],\n",
    ")\n",
    "ax.set_title(f\"Words written\")\n",
    "plt.savefig(figure_dir / \"total_words_written.pdf\", bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8465196358c50d9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for phase in phase_to_agent_to_word_count:\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(8, 6))\n",
    "    ax.pie(\n",
    "        phase_to_agent_to_word_count[phase].values(),\n",
    "        labels=phase_to_agent_to_word_count[phase].keys(),\n",
    "        autopct=\"%1.1f%%\",\n",
    "        colors=[agent_to_color[agent] for agent in phase_to_agent_to_word_count[phase]],\n",
    "    )\n",
    "    ax.set_title(f\"Words written in {phase.replace('_', ' ')}\")\n",
    "    plt.savefig(figure_dir / f\"{phase}_words_written.pdf\", bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ed39621",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da01ffda913de4b2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
